package com.llmops.demo;

import com.llmops.core.Context;
import com.llmops.core.Edge;
import com.llmops.core.WorkflowGraph;
import com.llmops.core.node.LLMNode;
import com.llmops.core.node.EndNode;
import com.llmops.core.node.StartNode;

import java.util.Collections;
import java.util.Map;

public class SubGraphTest {
    public static void main(String[] args) {
        WorkflowGraph workflow = new WorkflowGraph("mainWf");
        Context context = new Context();
        context.setVariable("sessionId",3);

        workflow.addNode(new StartNode("data_input","data_input",Collections.singletonMap("input","你是谁")));
        workflow.addNode(new LLMNode("llm_inference", "llm_inference",Map.of("systemPrompt","你是聊天助手回复用户的问题","userPrompt","{data_input_output}")));

        workflow.addNode(new EndNode("result_exporter", "result_exporter",Collections.singletonMap("content", "{sub_result_exporter_output}")));

        WorkflowGraph subWorkflow = new WorkflowGraph("mainWf","subWf");
        subWorkflow.addNode(new StartNode("sub_data_input","sub_data_input",Collections.singletonMap("input","{llm_inference_output}")));
        subWorkflow.addNode(new EndNode("sub_result_exporter","sub_result_exporter",Collections.singletonMap("content", "{sub_data_input_output}")));

        workflow.addNode(subWorkflow);

        workflow.addEdge(new Edge("data_input", "llm_inference"));
//        workflow.addEdge(new Edge("data_loader", "llm_inferenceB", Edge.ConditionType.EQUAL,"B"));

        workflow.addEdge(new Edge("llm_inference", "subWf"));

        workflow.addEdge(new Edge("subWf", "result_exporter"));
//        workflow.addEdge(new Edge("llm_inferenceB", "result_exporterB"));

        workflow.validate();
//        workflow.execute();

        workflow.execute(context);

        System.out.println("Workflow execution completed");
    }
}